from py3gpp import *
import numpy as np
import scipy.signal as signal
from NR.gscn import GSCN2CentFreq,gscn_pssIndices

from NR.pss import PSS

from Data.rawdata import RawData
# class RawData:
#     def __init__(self,fn) -> None:
#         self.fn = fn
#         self.handle = sigmf.sigmffile.fromfile(fn)

#         self.capture = self.handle.get_capture_info(0)
#         self.f_center = self.capture.get(sigmf.SigMFFile.FREQUENCY_KEY, 0)
#         self.fs = self.handle.get_global_field(sigmf.SigMFFile.SAMPLE_RATE_KEY)

#         self.f_start = self.f_center-self.fs/2
#         self.f_end = self.f_center+self.fs/2

#     def read(self,start=5e-3,end=30e-3):
#         if end>0:
#             self.waveform100 = self.handle.read_samples()[int(start * self.fs):int(end * self.fs)]
#         else:
#             self.waveform100 = self.handle.read_samples()[int(start * self.fs):]
#         self.waveform100 /= max(self.waveform100.real.max(), self.waveform100.imag.max())
            
#     def decimate(self,f):
#         df = -(f-self.f_center)/self.fs
#         self.waveform = signal.decimate(self.waveform100*np.exp(1j*2*np.pi*df*np.arange(len(self.waveform100))),self.dec_factor,ftype='fir')



# class NR(RawData):
#     def __init__(self,fn) -> None:
#         RawData.__init__(self,fn)

#         self.mu = 1
#         self.reset()
        
#     def reset(self):
#         self.carrier = nrCarrierConfig(NSizeGrid = 20, SubcarrierSpacing = 15 * 2**self.mu)
#         self.info = nrOFDMInfo(self.carrier)
#         self.Nfft = self.info['Nfft']

#         slotGrid = nrResourceGrid(self.carrier)
#         self.slotGrid = slotGrid[:, 0]
#         self.dec_factor = int(self.fs/(self.Nfft*15e3 * 2**self.mu))
        
#         self.detected_NID = -1
        
#         self.max_frames = 4
    
#     def set_search_end(self,up):
#         self.f_end = up
#         return self
    
#     def set_search_start(self,down):
#         self.f_start = down
#         return self
    
#     def decimate(self,gscn):
#         f_ssb = self.GSCN2CentFreq(gscn)
#         RawData.decimate(self,f_ssb)

#     def sync(self):
#         peak_value = np.zeros(3)
#         peak_index = np.zeros(3, 'int')
#         PSS_LEN = 128
#         NRE_PER_PRB = 12

#         start = 56
#         pssIndices = np.arange(start, start + PSS_LEN  - 1)
#         SampleRate = self.fs/8

#         for current_NID2 in np.arange(3, dtype='int'):
            
#             self.slotGrid[pssIndices] = nrPSS(current_NID2)
#             [refWaveform, info] = nrOFDMModulate(self.carrier, self.slotGrid, SampleRate = SampleRate)

#             refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP

#             temp = signal.correlate(self.waveform, refWaveform, 'valid')  # correlate over 25 ms
#             index = np.argmax(np.abs(temp))
#             peak = np.abs(temp[index])
#             peak_index[current_NID2] = index
#             peak_value[current_NID2] = peak

#         self.detected_NID2 = np.argmax(peak_value)
#         self.timingOffset = peak_index[self.detected_NID2]
#         self.peak =  peak_value[self.detected_NID2]
    
#     def GSCN2CentFreq(self,gscn):
#         self.gscn = gscn
#         if gscn<7499:
#             n = gscn % 3
#             if n==0:
#                 N = gscn // 3
#                 M = 3
#             elif n==1:
#                 N = gscn // 3
#                 M = 5
#             else:
#                 N = gscn // 3 + 1
#                 M = 1
#             return N*1.2e6+M*50e3
#         elif gscn<22256:
#             return (gscn-7499)*1.44e6+3e9
#         else:
#             return (gscn-22256)*17.28e6+24250.08e6
        
#     @staticmethod
#     def nearGSCN(fc):
#         if fc<3e9:
#             return int(round(fc/1.2e6*3))
#         elif fc<24.250e9:
#             return 7499+int(round((fc-3e9)/1.44e6))
#         else:
#             return 22256+int(round((fc-24250.08e6)/17.28))
        
#     def report(self):
#         r = f'{self.fn}:GSCN={self.gscn}  detect PSS {self.detected_NID2}@{self.timingOffset} with {self.peak}'
#         if self.detected_NID!=-1:
#             r += f' SSS {self.detected_NID}'
#         return r
    
#     def search(self,start=None,end=None):
#         if start==None:
#             start = self.f_start
#         if end==None:
#             end = self.f_end
#         sg = NR.nearGSCN(start)
#         se = NR.nearGSCN(end)
#         n = se - sg
#         offset = np.zeros((n,),dtype=np.int64)
#         nid2 = np.zeros((n,),dtype=np.int32)
#         peak = np.zeros((n,))
#         for k in range(n):
#             gscn = k+sg
#             self.decimate(gscn)
#             self.sync()
#             offset[k] = self.timingOffset
#             peak[k] = self.peak
#             nid2[k] = self.detected_NID2
#         k = np.argmax(peak)
#         self.gscn = k+sg
#         self.detected_NID2 = nid2[k]
#         self.peak = peak[k]
#         self.timingOffset = offset[k]
#         return k+sg,nid2[k],peak[k],offset[k]
    
#     def find_psses(self,waveform,alpha):
#         PSS_LEN = 128
#         NRE_PER_PRB = 12
#         start = 56
#         pssIndices = np.arange(start, start + PSS_LEN  - 1)
#         SampleRate = self.fs/8
#         all_pss = []
#         for nid2 in range(3):
#             self.slotGrid[pssIndices] = nrPSS(nid2)
#             [refWaveform, info] = nrOFDMModulate(self.carrier, self.slotGrid, SampleRate = SampleRate)
#             refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP
#             temp = signal.correlate(waveform, refWaveform, 'valid')  # correlate over 25 ms       
#             index = 10
#             while index < len(temp):
#                 if np.abs(temp[index])>self.peak*alpha:
#                     l_ = np.argmax(np.abs(temp[index-5:index+5]))
#                     all_pss.append((nid2, index-5+l_,temp[index-5+l_]))
#                     index += 10
#                 else:
#                     index += 1
#         return all_pss
    
#     def find_all_pss(self,alpha):
#         self.decimate(self.gscn)
#         self.all_pss = self.find_psses(self.waveform,alpha)
#         return self.all_pss
    
#     def sssCheck(self,time_off=None):
#         if time_off==None:
#             time_off = self.timingOffset
#         self.decimate(self.gscn)
#         tos = time_off - (127-118)*4
#         rxGrid = nrOFDMDemodulate(carrier = self.carrier, waveform = self.waveform[tos:], initialNSlot = 2, SampleRate=self.fs/self.dec_factor, CyclicPrefixFraction=0.5)
#         rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
#         sssEst = np.zeros(336*3)
#         start = 56
#         PSS_LEN = 128
#         sssIndices = np.arange(start, start + PSS_LEN  - 1)
#         sssRx = nrExtractResources(sssIndices, rxGrid[:,2:])
#         for ncellid in range(336*3):

#             sssRef = nrSSS(ncellid)
#             sssEst[ncellid] = np.abs(np.vdot(sssRx, sssRef))

#         self.detected_NID = np.argmax(sssEst)
#         return self.detected_NID
    
#     def detect_NID(self,waveform,nid2,offset):
#         scsSSB = 15 * 2**(self.mu)
#         nSlot = 0
#         rxSampleRate = self.fs/8
#         nrbSSB = 20

#         pss_pos = 592
#         rxGrid = nrOFDMDemodulate(waveform = waveform[offset-pss_pos:][:np.min((len(waveform), 2048*8))], nrb = nrbSSB, scs = scsSSB, initialNSlot = nSlot, SampleRate=rxSampleRate, CyclicPrefixFraction=0.5)
#         rxGrid = rxGrid[:,1:5]
#         rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
#         sssIndices = nrSSSIndices()
#         sssRx = nrExtractResources(sssIndices, rxGrid)
#         sssRx /= np.max((sssRx.real.max(), sssRx.imag.max())) # scale sssRx symbols individually
        
#         sssEst = np.zeros(336)
#         for NID1 in range(335):
#             ncellid = (3*NID1) + nid2
#             sssRef = nrSSS(ncellid)
#             sssEst[NID1] = np.abs(np.vdot(sssRx, sssRef))
#         detected_NID1 = np.argmax(sssEst)
#         peak = sssEst[detected_NID1]
#         detected_NID = detected_NID1*3 + nid2
#         return detected_NID,peak
        
#     def checkSSB(self,nid2,offset):
#         return self.checkSSBwithWaveform(self.waveform,nid2,offset)
    
#     def checkSSBwithWaveform(self,waveform,nid2,offset):
#         scsSSB = 15 * 2**(self.mu)
#         nSlot = 0
#         rxSampleRate = self.fs/8
#         nrbSSB = 20

#         pss_pos = 592
#         rxGrid = nrOFDMDemodulate(waveform = waveform[offset-pss_pos:][:np.min((len(waveform), 2048*8))], nrb = nrbSSB, scs = scsSSB, initialNSlot = nSlot, SampleRate=rxSampleRate, CyclicPrefixFraction=0.5)
#         rxGrid = rxGrid[:,1:5]
#         rxGrid /= np.max((rxGrid.real.max(), rxGrid.imag.max()))
#         sssIndices = nrSSSIndices()
#         sssRx = nrExtractResources(sssIndices, rxGrid)
#         sssRx /= np.max((sssRx.real.max(), sssRx.imag.max())) # scale sssRx symbols individually
        
#         sssEst = np.zeros(336)
#         for NID1 in range(335):
#             ncellid = (3*NID1) + nid2
#             sssRef = nrSSS(ncellid)
#             sssEst[NID1] = np.abs(np.vdot(sssRx, sssRef))
#         detected_NID1 = np.argmax(sssEst)
#         detected_NID = detected_NID1*3 + nid2
        
#         dmrsIndices = nrPBCHDMRSIndices(detected_NID, style='matlab')
#         xcorrPBCHDMRS = np.empty(7)
#         for ibar_SSB in range(7):
#             PBCHDMRS = nrPBCHDMRS(detected_NID, ibar_SSB)
#             xcorrPBCHDMRS[ibar_SSB] = np.abs(np.vdot(nrExtractResources(dmrsIndices, rxGrid), PBCHDMRS))
#         ibar_SSB = np.argmax(np.abs(xcorrPBCHDMRS))
#         # print(f"detect {detected_NID} ibar_ssb {ibar_SSB}")
        
#         refGrid = np.zeros((nrbSSB*12, 4), 'complex')
#         dmrsIndices = nrPBCHDMRSIndices(detected_NID)
#         temp = nrSetResources(dmrsIndices, refGrid, nrPBCHDMRS(detected_NID, ibar_SSB))
#         H, nVar = nrChannelEstimate(rxGrid = rxGrid, refGrid = refGrid)
#         pbchIndices = nrPBCHIndices(detected_NID)
#         pbch_eqed, csi = nrEqualizeMMSE(nrExtractResources(pbchIndices, rxGrid), nrExtractResources(pbchIndices, H), nVar)
#         pbchBits = nrSymbolDemodulate(pbch_eqed, 'QPSK', nVar, 'soft')
#         E = 864
#         v = ibar_SSB
#         scrambling_seq = nrPBCHPRBS(detected_NID, v, E)
#         scrambling_seq_bpsk = (-1)*scrambling_seq*2 + 1
#         pbchBits_descrambled = pbchBits * scrambling_seq_bpsk
#         pbchBits_csi = pbchBits_descrambled * np.repeat(csi, 2)
#         decoded3, crc_result, payload, lsbotfsfn, hrf, msbidxoffset = nrBCHDecode(pbchBits_csi,8,0,detected_NID)
         
#         return crc_result, detected_NID, lsbotfsfn, ibar_SSB
        
#     def checkSSB_all(self):
#         self.all_ssb = []
#         print(f"checking SSB, total: {len(self.all_pss)}")
#         for nid2,offset,_ in self.all_pss:
#             crc_result, detected_NID, lsbotfsfn, _, _ = self.checkSSB(nid2,offset)
#             if crc_result==0:
#                 self.all_ssb.append(SSBInfo().cellid(detected_NID).offset(offset-592))
#         return self.all_ssb
    
#     def step0(self):
#         self.read(end=30e-3)
#         gscn,nid2,peak,offset = self.search()
#         self.read(end=-1.0)
#         self.decimate(gscn)
#         self.frames = int(len(self.waveform)/(self.fs/8*20e-3))-1
#         crc_result, detected_NID, _, ibar_SSB = self.checkSSB(self.detected_NID2,self.timingOffset) 
#         if crc_result==0:
#             self.gscn = gscn
#             self.detected_NID = detected_NID
#             self.ibar_SSB = ibar_SSB
#             self.synced = True
#         else:
#             self.synced = False
#         return crc_result==0
    
#     def report_nid(self):
#         if self.synced:
#             return f"gscn={self.gscn}, nid={self.detected_NID}, ibar={self.ibar_SSB}"
#         else:
#             return "no synced"
    
#     def search_t(self, wv, ref):
#         l = len(wv)
#         res = 256
#         step = 1/res
#         fwv = np.fft.fft(wv)
#         fref = np.fft.fft(ref)
#         est = np.zeros((4*res,))
#         for i in range(4*res):
#             est[i] = np.abs(np.dot(fwv*np.exp(1j*2*np.pi*np.arange(l)*(i-2*res)*step/l),fref.conj()))
#         i = np.argmax(est)
#         return (i-2*res)*step, est
    
#     def calibration_t(self, wv, t):
#         l = len(wv)
#         fwv = np.fft.fft(wv)
#         ret = np.fft.ifft(fwv*np.exp(1j*2*np.pi*np.arange(l)*t/l))
#         return ret
    
#     def est(self):
#         SampleRate = self.fs/8
#         ss = int(round(20e-3*SampleRate))
#         slotGrid = nrResourceGrid(self.carrier)
#         slotGrid = slotGrid[:,:4]
#         nrSetResources(nrPSSIndices(),slotGrid,nrPSS(self.detected_NID%3))
#         nrSetResources(nrSSSIndices(),slotGrid,nrSSS(self.detected_NID))
#         dmrsIndices = nrPBCHDMRSIndices(self.detected_NID)
#         nrSetResources(dmrsIndices, slotGrid, nrPBCHDMRS(self.detected_NID, self.ibar_SSB))
#         [refWaveform, info] = nrOFDMModulate(self.carrier, slotGrid, SampleRate = SampleRate)
#         refWaveform = refWaveform[info['CyclicPrefixLengths'][0]:]; # remove CP
#         ref_len = len(refWaveform)
#         self.phase = np.zeros((self.frames,),dtype=np.complex128)
#         self.dt = np.zeros((self.frames,))
#         for i in range(self.frames):
#             wv = self.waveform[self.timingOffset+i*ss:self.timingOffset+i*ss+ref_len]
#             self.dt[i],_ = self.search_t(wv,refWaveform)
#             c_wv = self.calibration_t(wv,self.dt[i])
#             self.phase[i] = np.dot(c_wv,refWaveform.conj())

#     def step1(self,alpha=1.0):
#         SampleRate = self.fs/8
#         ss = int(round(20e-3*SampleRate))
#         self.est()
#         ssb_len = 512*8
#         self.ssbWaveform = np.zeros((2*ssb_len,),dtype=np.complex64)
#         for i in range(min(self.frames,self.max_frames)):
#             wv = self.waveform[self.timingOffset+i*ss-ssb_len:self.timingOffset+i*ss+ssb_len]
#             c_wv = self.calibration_t(wv,self.dt[i])
#             self.ssbWaveform += c_wv * self.phase[i].conj()
#         self.ssbWaveform /= max(self.ssbWaveform.real.max(), self.ssbWaveform.imag.max()) 
#         self.step1_pss = self.find_psses(self.ssbWaveform,alpha) 
#         self.step1_pss.sort(key=lambda x:np.abs(x[2]))
#         return self.step1_pss
    
#     def step2(self):
#         ssb_len = 512*8
#         self.ssbs = []
#         self.step2_sss = []
#         for nid2, off, _ in self.step1_pss[:10]:
#             crc_result, detected_NID, _, ibar_SSB = self.checkSSB(nid2,self.timingOffset+off-ssb_len)
#             detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
#             self.step2_sss.append((detected_NID,np.abs(peak)))
        
#     def step3(self):
#         ssb_len = 512*8
#         self.ssbs = []
#         for nid2, off, _ in self.step1_pss[:10]:
#             crc_result, detected_NID, _, ibar_SSB = self.checkSSB(nid2,self.timingOffset+off-ssb_len)
#             detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
#             if crc_result == 0:
#                 detected_NID, peak = self.detect_NID(self.ssbWaveform,nid2,off)
#                 self.ssbs.append(SSBInfo().cellid(detected_NID).offset(self.timingOffset+off-ssb_len).set_peak(peak).set_ibar(ibar_SSB))
#         return self.ssbs
          
                    
class SSBInfo:
    def __init__(self):
        pass
    def set_init_slot(self,i):
        self.initialNSlot = i
        return self
    def cellid(self, nid):
        self.nid = nid
        return self
    def offset(self,off):
        self.off = off
        return self
    def set_ibar(self,i):
        self.ibar = i
        return self
    def set_peak(self,peak):
        self.peak = peak
        return self
    def __str__(self):
        return f"NID {self.nid} @{self.off} peak={self.peak} ibar={self.ibar}"
    
    def ref3(self,carrier,fs):
        slotGrid = nrResourceGrid(carrier)
        slotGrid = slotGrid[:,:4]
        nrSetResources(nrPSSIndices(),slotGrid,nrPSS(self.nid%3))
        nrSetResources(nrSSSIndices(),slotGrid,nrSSS(self.nid))
        dmrsIndices = nrPBCHDMRSIndices(self.nid)
        nrSetResources(dmrsIndices, slotGrid, nrPBCHDMRS(self.nid, self.ibar))
        [refWaveform, info] = nrOFDMModulate(carrier, slotGrid, SampleRate = fs, initialNSlot=self.initialNSlot)
        return refWaveform
    
    def ref2(self,carrier,fs):
        slotGrid = nrResourceGrid(carrier)
        slotGrid = slotGrid[:,:4]
        nrSetResources(nrPSSIndices(),slotGrid,nrPSS(self.nid%3))
        nrSetResources(nrSSSIndices(),slotGrid,nrSSS(self.nid))
        [refWaveform, info] = nrOFDMModulate(carrier, slotGrid, SampleRate = fs)
        return refWaveform
    
    def ref1(self,carrier,fs):
        slotGrid = nrResourceGrid(carrier)
        slotGrid = slotGrid[:,:4]
        nrSetResources(nrPSSIndices(),slotGrid,nrPSS(self.nid%3))
        [refWaveform, info] = nrOFDMModulate(carrier, slotGrid, SampleRate = fs)
        return refWaveform
    
    def ref_sss(self,carrier,fs):
        slotGrid = nrResourceGrid(carrier)
        slotGrid = slotGrid[:,:4]
        nrSetResources(nrSSSIndices(),slotGrid,nrSSS(self.nid))
        [refWaveform, info] = nrOFDMModulate(carrier, slotGrid, SampleRate = fs)
        return refWaveform
    
    def ref_pbch_dmrs(self,carrier,fs):
        slotGrid = nrResourceGrid(carrier)
        slotGrid = slotGrid[:,:4]
        dmrsIndices = nrPBCHDMRSIndices(self.nid)
        nrSetResources(dmrsIndices, slotGrid, nrPBCHDMRS(self.nid, self.ibar))
        [refWaveform, info] = nrOFDMModulate(carrier, slotGrid, SampleRate = fs)
        return refWaveform